import matplotlib.pyplot as plt
import pandas as pd
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score


def score_clustering(df):
    """
    :param df:
    :return: 合理的聚类结果和可视化图
    """

    # 提取数值型列并转换为矩阵
    X = df.select_dtypes(include='number').values

    # 计算聚类数量和最佳轮廓系数
    best_n_clusters, best_score = 3, -1
    for n_clusters in range(3, min(8, X.shape[1] + 1)):
        kmeans = KMeans(n_clusters=n_clusters, n_init=10)
        labels = kmeans.fit_predict(X)
        score = silhouette_score(X, labels)
        if score > best_score:
            best_n_clusters = n_clusters
            best_score = score

    # 根据数据列数选择绘制二维或三维图像
    plt.rcParams['font.family'] = ['STSong']
    if X.shape[1] == 2:
        # 定义和初始化 kmeans
        kmeans = KMeans(n_clusters=best_n_clusters, n_init=10)
        labels = kmeans.fit_predict(X)

        results = []
        for i in range(best_n_clusters):
            cluster_indices = df.index[labels == i].tolist()
            results.append(cluster_indices)
        print(results)
        for i in range(len(results)):
            print("第" + str(i + 1) + "类学生包括：")
            for j in range(len(results[i])):
                print(df.iloc[j, 0], end='')
                print("、", end='')

        plt.figure()
        for i in range(best_n_clusters):
            plt.scatter(X[labels == i, 0], X[labels == i, 1], s=50, label='Cluster {}'.format(i + 1))
        plt.scatter(kmeans.cluster_centers_[:, 0], kmeans.cluster_centers_[:, 1], s=150, marker='*', color='black',
                    label='Centroids')
        plt.title('学生聚类结果，类数为{}'.format(best_n_clusters))
        plt.xlabel(df.select_dtypes(include='number').columns[0])
        plt.ylabel(df.select_dtypes(include='number').columns[1])
        plt.legend()
        plt.show()

    elif X.shape[1] >= 3:
        # 定义和初始化 kmeans
        kmeans = KMeans(n_clusters=best_n_clusters, n_init=10)
        labels = kmeans.fit_predict(X)

        results = []
        for i in range(best_n_clusters):
            cluster_indices = df.index[labels == i].tolist()
            results.append(cluster_indices)
        print(results)
        for i in range(len(results)):
            print("第" + str(i + 1) + "类学生包括：")
            for j in results[i]:
                print(df.iloc[j, 0], end='')
                print("、", end='')

        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        for i in range(best_n_clusters):
            ax.scatter(X[labels == i, 0], X[labels == i, 1], X[labels == i, 2], s=50, label='Cluster {}'.format(i + 1))
        ax.scatter(kmeans.cluster_centers_[:, 0], kmeans.cluster_centers_[:, 1], kmeans.cluster_centers_[:, 2],
                   s=150, marker='*', color='black', label='Centroids')
        ax.set_title('学生聚类结果，类数为{}'.format(best_n_clusters))
        ax.set_xlabel(df.select_dtypes(include='number').columns[0])
        ax.set_ylabel(df.select_dtypes(include='number').columns[1])
        ax.set_zlabel(df.select_dtypes(include='number').columns[2] if X.shape[1] >= 3 else '')
        ax.legend()
        plt.show()

    # df.to_excel('test_score_result001.xlsx', index=None)


def main():
    df = pd.read_excel("../TestExample/test_score001.xlsx")
    score_clustering(df)


if __name__ == "__main__":
    main()
